import cv2


def seg_result_proc(src_img_data,dst_img_path,results,color_dic):
    ret = {"clses":[],"points":[]}
    clses = results.boxes.cls
    xyz = results.masks.xy
    for index,points in enumerate(xyz):
        if len(points)>0:
            cls = int(clses[index].item())
            ret["clses"].append(cls)
            ret["points"].append(points.astype(int).tolist())
            color = color_dic["data"][cls]
            cv2.drawContours(src_img_data,[points.astype(int)],-1,color,2)
            cv2.imwrite(dst_img_path,src_img_data)
    ret

def obt_result_proc(src_img_data,dst_img_path,results,color_dic):
    ret = {"clses":[],"points":[],"conf":[]}
    clses = results.boxes.cls
    confidences = results.boxes.conf
    xyz = results.boxes.xyxy
    names = results.names
    for index,points in enumerate(xyz):
        cls = int(clses[index].item())
        conf = "{:.2f}".format(float(confidences[index].item()))
        cls_name = names[cls]
        ret["clses"].append(cls)
        points_np = points.numpy().astype(int)
        x1,y1,x2,y2 = points_np[0],points_np[1],points_np[2],points_np[3]
        ret["points"].append(points_np.tolist())
        color = color_dic["data"][cls]
        cv2.rectangle(src_img_data,(x1,y1),(x2,y2),color,2)
        text = f"{cls_name}-{conf}"
        cv2.putText(src_img_data,text,(x1,y1-10),cv2.FONT_HERSHEY_SIMPLEX,0.5,color,2)
        cv2.imwrite(dst_img_path,src_img_data)
    ret
